/**
 * @ProjectName: 
 * @Copyright: 2014 lisheng  All Right Reserved.
 * @address: toughheart@163.com
 * @date: 2015年1月2日 下午8:21:18
 * @Description: 本内容未经本人允许禁止使用、转发.
 */

package org.apache.lucene.demo;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;

import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.facet.FacetResult;
import org.apache.lucene.facet.Facets;
import org.apache.lucene.facet.FacetsCollector;
import org.apache.lucene.facet.FacetsConfig;
import org.apache.lucene.facet.LabelAndValue;
import org.apache.lucene.facet.range.LongRange;
import org.apache.lucene.facet.range.LongRangeFacetCounts;
import org.apache.lucene.facet.taxonomy.CachedOrdinalsReader;
import org.apache.lucene.facet.taxonomy.DocValuesOrdinalsReader;
import org.apache.lucene.facet.taxonomy.FacetLabel;
import org.apache.lucene.facet.taxonomy.FastTaxonomyFacetCounts;
import org.apache.lucene.facet.taxonomy.OrdinalsReader;
import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager;
import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager.SearcherAndTaxonomy;
import org.apache.lucene.facet.taxonomy.TaxonomyFacetCounts;
import org.apache.lucene.facet.taxonomy.TaxonomyFacetSumValueSource;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyWriter;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig.OpenMode;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.valuesource.BytesRefFieldSource;
import org.apache.lucene.queries.function.valuesource.IntFieldSource;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.grouping.AbstractAllGroupsCollector;
import org.apache.lucene.search.grouping.AbstractFirstPassGroupingCollector;
import org.apache.lucene.search.grouping.AbstractGroupFacetCollector;
import org.apache.lucene.search.grouping.AbstractSecondPassGroupingCollector;
import org.apache.lucene.search.grouping.GroupingSearch;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.function.FunctionAllGroupsCollector;
import org.apache.lucene.search.grouping.function.FunctionFirstPassGroupingCollector;
import org.apache.lucene.search.grouping.function.FunctionSecondPassGroupingCollector;
import org.apache.lucene.search.grouping.term.TermAllGroupsCollector;
import org.apache.lucene.search.grouping.term.TermFirstPassGroupingCollector;
import org.apache.lucene.search.grouping.term.TermGroupFacetCollector;
import org.apache.lucene.search.grouping.term.TermSecondPassGroupingCollector;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.mutable.MutableValue;
import org.apache.lucene.util.mutable.MutableValueStr;

/**
 * 
 * @author lisheng
 * @date 2015年1月2日 下午8:21:18
 * @version V1.0
 */
public class LuceneFacetServerImpl extends LuceneServerImpl {

	protected SearcherTaxonomyManager taxonomyManager;

	protected DirectoryTaxonomyWriter taxonomyWriter;

	protected Directory taxonomyDir = null;

	public LuceneFacetServerImpl() {
		super();
	}

	public LuceneFacetServerImpl(String address) {
		super(address);
	}

	public LuceneFacetServerImpl(Directory dir) {
		super(dir);
	}

	public LuceneFacetServerImpl(IndexWriter indexWriter,
			IndexReader indexReader) {
		super(indexWriter, indexReader);
	}

	public void addCategory(String... name) throws IOException {
		taxonomyWriter.addCategory(new FacetLabel(name));
	}

	public void openTaxonomy(String address, OpenMode openMode)
			throws IOException {
		long start = System.currentTimeMillis();
		this.logger.info("打开服务" + address + "....");
		if (taxonomyDir == null && StringUtils.isNotBlank(address)) {
			taxonomyDir = this.openDir(address);
		}
		if (this.taxonomyWriter == null && taxonomyDir != null) {
			this.taxonomyWriter = new DirectoryTaxonomyWriter(taxonomyDir,
					openMode);
		}
		if (this.nearRealTime) {
			taxonomyManager = new SearcherTaxonomyManager(this.indexWriter,
					true, this.factory, this.taxonomyWriter);
		} else {
			taxonomyManager = new SearcherTaxonomyManager(dir, taxonomyDir,
					this.factory);
		}
		logger.info("打开服务" + address + "耗时:"
				+ (System.currentTimeMillis() - start));
	}

	public Facets searchFacet(Query query) throws IOException {
		Facets facets = null;
		facets = this.search(query, new FacetCollectorMapper<Facets>() {

			@Override
			public Facets mapper(SearcherAndTaxonomy pair,
					FacetsCollector collector) {
				try {
					return getTaxonomyFacetCounts(pair.taxonomyReader, config,
							collector);
				} catch (IOException e) {
					e.printStackTrace();
				}
				return null;
			}
		});
		return facets;
	}

	public <T> T search(Query query, FacetCollectorMapper<T> mapper)
			throws IOException {
		SearcherAndTaxonomy pair = this.taxonomyManager.acquire();
		FacetsCollector sfc = new FacetsCollector();
		System.out.println(query);
		T t = null;
		try {
			pair.searcher.search(query, sfc);
			t = mapper.mapper(pair, sfc);
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			try {
				taxonomyManager.release(pair);
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		return t;
	}

	public interface CollectorMapper<T> {
		public T mapper(Collector collector);
	}

	public interface FacetCollectorMapper<T> {
		public T mapper(SearcherAndTaxonomy pair, FacetsCollector collector);
	}

	public TaxonomyFacetSumValueSource sum(Query query, final String fieldName)
			throws IOException {
		TaxonomyFacetSumValueSource facets = null;
		facets = this.search(query,
				new FacetCollectorMapper<TaxonomyFacetSumValueSource>() {

					@Override
					public TaxonomyFacetSumValueSource mapper(
							SearcherAndTaxonomy pair, FacetsCollector collector) {
						try {
							return new TaxonomyFacetSumValueSource(
									pair.taxonomyReader, config, collector,
									new IntFieldSource(fieldName));
						} catch (IOException e) {
							e.printStackTrace();
						}
						return null;
					}
				});
		return facets;
	}

	public Facets rangeCount(Query query, final String fieldName,
			final LongRange... ranges) throws IOException {
		Facets facets = null;
		facets = this.search(query, new FacetCollectorMapper<Facets>() {

			@Override
			public Facets mapper(SearcherAndTaxonomy pair,
					FacetsCollector collector) {
				try {
					return new LongRangeFacetCounts(fieldName, collector,
							ranges);
				} catch (IOException e) {
					e.printStackTrace();
				}
				return null;
			}
		});
		return facets;
	}

	public void replaceTaxonomy(Directory taxoDir) throws IOException {
		taxonomyWriter.replaceTaxonomy(taxoDir);
		taxonomyManager.maybeRefresh();
	}

	protected AbstractGroupFacetCollector createCollector(String groupField,
			String facetField, String facetPrefix,
			boolean multipleFacetsPerDocument, int initSize) {
		BytesRef facetPrefixBR = facetPrefix == null ? null : new BytesRef(
				facetPrefix);
		// DocValues cannot be multi-valued:
		assert !multipleFacetsPerDocument || !groupField.endsWith("_dv");
		return TermGroupFacetCollector.createTermGroupFacetCollector(
				groupField, facetField, multipleFacetsPerDocument,
				facetPrefixBR, initSize);
	}

	public void group() {
		// GroupingSearch groupingSearch = null;
		// TopGroups<?> groups = groupingSearch.search();

	}

	protected GroupingSearch createRandomGroupingSearch(String groupField,
			boolean isFunction, Sort groupSort, int docsInGroup, boolean cache) {
		GroupingSearch groupingSearch;
		if (isFunction) {
			ValueSource vs = new BytesRefFieldSource(groupField);
			groupingSearch = new GroupingSearch(vs, new HashMap<>());
		} else {
			groupingSearch = new GroupingSearch(groupField);
		}

		groupingSearch.setGroupSort(groupSort);
		groupingSearch.setGroupDocsLimit(docsInGroup);

		if (cache) {
			groupingSearch.setCachingInMB(4.0, true);
		}

		return groupingSearch;
	}

	protected AbstractFirstPassGroupingCollector<?> createRandomFirstPassCollector(
			String groupField, boolean isFunction, Sort groupSort, int topDocs)
			throws IOException {
		AbstractFirstPassGroupingCollector<?> selected;
		if (isFunction) {
			ValueSource vs = new BytesRefFieldSource(groupField);
			selected = new FunctionFirstPassGroupingCollector(vs,
					new HashMap<>(), groupSort, topDocs);
		} else {
			selected = new TermFirstPassGroupingCollector(groupField,
					groupSort, topDocs);
		}
		return selected;
	}

	protected AbstractFirstPassGroupingCollector<?> createFirstPassCollector(
			String groupField, Sort groupSort, int topDocs,
			AbstractFirstPassGroupingCollector<?> firstPassGroupingCollector)
			throws IOException {
		if (TermFirstPassGroupingCollector.class
				.isAssignableFrom(firstPassGroupingCollector.getClass())) {
			ValueSource vs = new BytesRefFieldSource(groupField);
			return new FunctionFirstPassGroupingCollector(vs, new HashMap<>(),
					groupSort, topDocs);
		} else {
			return new TermFirstPassGroupingCollector(groupField, groupSort,
					topDocs);
		}
	}

	@SuppressWarnings({ "unchecked", "rawtypes" })
	protected <T> AbstractSecondPassGroupingCollector<T> createSecondPassCollector(
			AbstractFirstPassGroupingCollector firstPassGroupingCollector,
			String groupField, Sort groupSort, Sort sortWithinGroup,
			int groupOffset, int maxDocsPerGroup, boolean getScores,
			boolean getMaxScores, boolean fillSortFields) throws IOException {

		if (TermFirstPassGroupingCollector.class
				.isAssignableFrom(firstPassGroupingCollector.getClass())) {
			Collection<SearchGroup<BytesRef>> searchGroups = firstPassGroupingCollector
					.getTopGroups(groupOffset, fillSortFields);
			return (AbstractSecondPassGroupingCollector) new TermSecondPassGroupingCollector(
					groupField, searchGroups, groupSort, sortWithinGroup,
					maxDocsPerGroup, getScores, getMaxScores, fillSortFields);
		} else {
			ValueSource vs = new BytesRefFieldSource(groupField);
			Collection<SearchGroup<MutableValue>> searchGroups = firstPassGroupingCollector
					.getTopGroups(groupOffset, fillSortFields);
			return (AbstractSecondPassGroupingCollector) new FunctionSecondPassGroupingCollector(
					searchGroups, groupSort, sortWithinGroup, maxDocsPerGroup,
					getScores, getMaxScores, fillSortFields, vs, new HashMap());
		}
	}

	// Basically converts searchGroups from MutableValue to BytesRef if grouping
	// by ValueSource
	protected AbstractSecondPassGroupingCollector<?> createSecondPassCollector(
			AbstractFirstPassGroupingCollector<?> firstPassGroupingCollector,
			String groupField, Collection<SearchGroup<BytesRef>> searchGroups,
			Sort groupSort, Sort sortWithinGroup, int maxDocsPerGroup,
			boolean getScores, boolean getMaxScores, boolean fillSortFields)
			throws IOException {
		if (firstPassGroupingCollector.getClass().isAssignableFrom(
				TermFirstPassGroupingCollector.class)) {
			return new TermSecondPassGroupingCollector(groupField,
					searchGroups, groupSort, sortWithinGroup, maxDocsPerGroup,
					getScores, getMaxScores, fillSortFields);
		} else {
			ValueSource vs = new BytesRefFieldSource(groupField);
			List<SearchGroup<MutableValue>> mvalSearchGroups = new ArrayList<>(
					searchGroups.size());
			for (SearchGroup<BytesRef> mergedTopGroup : searchGroups) {
				SearchGroup<MutableValue> sg = new SearchGroup<>();
				MutableValueStr groupValue = new MutableValueStr();
				if (mergedTopGroup.groupValue != null) {
					groupValue.value.copyBytes(mergedTopGroup.groupValue);
				} else {
					groupValue.exists = false;
				}
				sg.groupValue = groupValue;
				sg.sortValues = mergedTopGroup.sortValues;
				mvalSearchGroups.add(sg);
			}

			return new FunctionSecondPassGroupingCollector(mvalSearchGroups,
					groupSort, sortWithinGroup, maxDocsPerGroup, getScores,
					getMaxScores, fillSortFields, vs, new HashMap<>());
		}
	}

	protected AbstractAllGroupsCollector<?> createAllGroupsCollector(
			AbstractFirstPassGroupingCollector<?> firstPassGroupingCollector,
			String groupField) {
		if (firstPassGroupingCollector.getClass().isAssignableFrom(
				TermFirstPassGroupingCollector.class)) {
			return new TermAllGroupsCollector(groupField);
		} else {
			ValueSource vs = new BytesRefFieldSource(groupField);
			return new FunctionAllGroupsCollector(vs, new HashMap<>());
		}
	}

	public Facets getTaxonomyFacetCounts(TaxonomyReader taxoReader,
			FacetsConfig config, FacetsCollector c) throws IOException {
		return getTaxonomyFacetCounts(taxoReader, config, c,
				FacetsConfig.DEFAULT_INDEX_FIELD_NAME);
	}

	public Facets getTaxonomyFacetCounts(TaxonomyReader taxoReader,
			FacetsConfig config, FacetsCollector c, String indexFieldName)
			throws IOException {

		return this.getTaxonomyFacetCounts(taxoReader, config, c,
				indexFieldName, true, false);
	}

	public Facets getTaxonomyFacetCounts(TaxonomyReader taxoReader,
			FacetsConfig config, FacetsCollector c, String indexFieldName,
			boolean fast, boolean cache) throws IOException {
		Facets facets;
		if (fast) {
			facets = new FastTaxonomyFacetCounts(indexFieldName, taxoReader,
					config, c);
		} else {
			OrdinalsReader ordsReader = new DocValuesOrdinalsReader(
					indexFieldName);
			if (cache) {
				ordsReader = new CachedOrdinalsReader(ordsReader);
			}
			facets = new TaxonomyFacetCounts(ordsReader, taxoReader, config, c);
		}

		return facets;
	}

	protected void sortTies(List<FacetResult> results) {
		for (FacetResult result : results) {
			sortTies(result.labelValues);
		}
	}

	protected void sortTies(LabelAndValue[] labelValues) {
		double lastValue = -1;
		int numInRow = 0;
		int i = 0;
		while (i <= labelValues.length) {
			if (i < labelValues.length
					&& labelValues[i].value.doubleValue() == lastValue) {
				numInRow++;
			} else {
				if (numInRow > 1) {
					Arrays.sort(labelValues, i - numInRow, i,
							new Comparator<LabelAndValue>() {
								@Override
								public int compare(LabelAndValue a,
										LabelAndValue b) {
									assert a.value.doubleValue() == b.value
											.doubleValue();
									return new BytesRef(a.label)
											.compareTo(new BytesRef(b.label));
								}
							});
				}
				numInRow = 1;
				if (i < labelValues.length) {
					lastValue = labelValues[i].value.doubleValue();
				}
			}
			i++;
		}
	}

	protected void sortLabelValues(List<LabelAndValue> labelValues) {
		Collections.sort(labelValues, new Comparator<LabelAndValue>() {
			@Override
			public int compare(LabelAndValue a, LabelAndValue b) {
				if (a.value.doubleValue() > b.value.doubleValue()) {
					return -1;
				} else if (a.value.doubleValue() < b.value.doubleValue()) {
					return 1;
				} else {
					return new BytesRef(a.label)
							.compareTo(new BytesRef(b.label));
				}
			}
		});
	}

	protected void sortFacetResults(List<FacetResult> results) {
		Collections.sort(results, new Comparator<FacetResult>() {
			@Override
			public int compare(FacetResult a, FacetResult b) {
				if (a.value.doubleValue() > b.value.doubleValue()) {
					return -1;
				} else if (b.value.doubleValue() > a.value.doubleValue()) {
					return 1;
				} else {
					return 0;
				}
			}
		});
	}

	@Override
	public void close() throws IOException {
		IOUtils.close(taxonomyManager, taxonomyWriter, indexWriter,
				taxonomyDir, dir);
	}

	@Override
	public boolean open(String address, Analyzer analyzer, OpenMode openMode) {
		boolean flag = super.open(address, analyzer, openMode);
		File f = new File(address);

		if (f.exists()) {
			String taxoDir = f.getParent() + File.separator + "taxoIndex";
			try {
				this.openTaxonomy(taxoDir, OpenMode.CREATE_OR_APPEND);
				flag = true;
			} catch (IOException e) {
				e.printStackTrace();
			}
		} else {
			flag = false;
		}
		return flag;
	}

	FacetsConfig config = new FacetsConfig();

	@Override
	public void addDoc(Document doc) {
		System.out.println(doc);
		System.out.println();
		Document doc1 = null;
		try {
			doc1 = config.build(this.taxonomyWriter, doc);
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println(doc1);
		super.addDoc(doc1);
	}

	@Override
	public void deleteAllDoc() {
		super.deleteAllDoc();
	}

	@Override
	public void commit() {
		super.commit();
		if (taxonomyWriter != null) {
			try {
				this.taxonomyWriter.commit();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}
}
